from torchvision.datasets import CIFAR10, CIFAR100
from torchvision.transforms import transforms
from torch.utils.data import DataLoader, Dataset
import torchvision
from .mocov2 import get_mocov2_loader
from .simsiam import get_simsiam_loader
from .mae import get_mae_loader
from .densecl import get_densecl_loader
from .aug import test_transform,train_transform
from .simclr import get_simclr_loader

def get_loader(args,method="mocov2"):
    """
    获取数据集
    """
    if method == 'mocov2':
        return get_mocov2_loader(args)
    elif method =='simsiam':
        return get_simsiam_loader(args)
    elif method =='memory':
        # 使用knn 进行预训练测试
        return get_memory_loader(args)
    elif method =='supervised':
        #  有监督的预训练
        return get_supervised_loader(args)
    elif method =='mae':
        return get_mae_loader(args)
    elif method=="densecl":
        return get_densecl_loader(args)
    elif method=="simclr":
        return get_simclr_loader(args)
    else:
        raise NotImplementedError

def get_memory_loader(args):

    if args.dataset == 'cifar10':
        trainset = torchvision.datasets.CIFAR10(root=args.data_path, train=True, download=True,
                                                transform=test_transform(args))
        memory_loader = DataLoader(
            trainset, batch_size=args.batch_size, drop_last=True, shuffle=False, num_workers=args.num_workers)
        return memory_loader

    elif args.dataset == 'cifar100':
        trainset = torchvision.datasets.CIFAR100(root=args.data_path, train=True, download=True,
                                                 transform=test_transform(args))

        memory_loader = DataLoader(
            trainset, batch_size=args.batch_size, drop_last=True, shuffle=False, num_workers=args.num_workers)

        return memory_loader
    else:
        raise NotImplementedError


def get_supervised_loader(args):
    """
    有监督训练的 loader
    """
    if args.dataset == 'cifar10':
        trainset = torchvision.datasets.CIFAR10(root=args.data_path, train=True, download=True,
                                                transform=train_transform(args))
        testset = torchvision.datasets.CIFAR10(root=args.data_path, train=False, download=True,
                                               transform=test_transform(args))

        trainloader = DataLoader(
            trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers)
        testloader = DataLoader(
            testset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
        return trainloader, testloader

    elif args.dataset == 'cifar100':
        trainset = torchvision.datasets.CIFAR100(root=args.data_path, train=True, download=True,
                                                 transform=train_transform(args))
        testset = torchvision.datasets.CIFAR100(root=args.data_path, train=False, download=True,
                                                transform=test_transform(args))
        trainloader = DataLoader(
            trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers)
        testloader = DataLoader(
            testset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)

        return trainloader, testloader
    else:
        raise NotImplementedError